# deepseek
import os
import argparse
import sys
from locust import HttpUser, task, between, events
from config.all_paths import get_root_path
from config.global_variable import ENV_DICT
from utils.file_handling.excel_handling import ExcelUtil
from utils.interface_auth.auth_tool import Authentication
from utils.performance_tool.config_loader import load_config
from utils.performance_tool.load_shapes import SHAPE_MAP
from utils.request_tool.send_request import RequestSender
from utils.time_tool.time_control import parse_duration

# 解析命令行参数
def parse_custom_args():
    """解析自定义命令行参数"""
    parser = argparse.ArgumentParser()
    parser.add_argument('--env', required=True, help='目标环境名称')
    parser.add_argument('--role', required=True, help='测试角色类型')
    parser.add_argument('--profile_name', required=True, help='测试场景名称')
    args, remaining = parser.parse_known_args()
    # 过滤已解析的自定义参数，避免干扰Locust
    sys.argv = sys.argv[:1] + remaining  
    return args

# 获取命令行参数
args = parse_custom_args()
ENV = args.env
ROLE = args.role
PROFILE_NAME = args.profile_name

# 加载配置
config = load_config(env=ENV, role=ROLE)
profile_config = config['profiles'].get(PROFILE_NAME, {})

# 数据准备
excelName = 'data/case_excel/wutong/uat/医生账号相关接口.xlsx'
testData = ExcelUtil(os.path.join(get_root_path(), excelName)).read_excel()

@events.test_start.add_listener
def init_test(environment, **kwargs):
    """测试初始化"""
    print(f"\n=== 启动 {PROFILE_NAME} 测试 ===")
    print(f"目标环境: {ENV}")
    print(f"测试角色: {ROLE}")
    print(f"基准URL: {config['env'].get('url')}")

@events.test_stop.add_listener
def finalize_test(environment, **kwargs):
    """测试收尾"""
    print("\n=== 测试结果摘要 ===")
    print(f"总请求数: {environment.stats.total.num_requests}")
    print(f"失败率: {environment.stats.total.fail_ratio:.2%}")

class BaseUser(HttpUser):
    wait_time = between(0.5, 2)
    abstract = True

    def on_start(self):
        """认证初始化"""
        env_config = ENV_DICT.get(ENV, {})
        self.s = Authentication(
            username=env_config.get('username'),
            password=env_config.get('password'),
            base_url=env_config.get('url'),
        ).auth_wutong()

class BaselineUser(BaseUser):
    @task
    def health_check(self):
        RequestSender().send_request(self.s, testData[0])

class StressUser(BaseUser):
    @task(5)
    def create_data(self):
        RequestSender().send_request(self.s, testData[0])

class ConcurrencyUser(BaseUser):
    @task
    def update_resource(self):
        resource_id = profile_config.get("concurrency_key")
        self.client.put(f"/resources/{resource_id}")

class WaveUser(BaseUser):
    @task
    def wave_task(self):
        self.client.get("/api/wave-test")

# 动态测试类
class DynamicTestClass(eval(profile_config.get("user_class", "BaselineUser"))):
    pass

# 负载形状配置
class CustomLoadShape(SHAPE_MAP[profile_config.get("load_shape", "constant")]):
    shape_users = profile_config.get("users", 1)
    spawn_rate = profile_config.get("spawn_rate", 1)
    shape_duration = parse_duration(profile_config.get("duration", "1m"))
    shape_stages = profile_config.get("stages", [])

# 暴露给Locust的配置
class TestScenario(DynamicTestClass):
    pass

def get_shape_class():
    return CustomLoadShape